import math
import re
import string

import gradio as gr
import mindspore
import mindspore.dataset as dataset
import mindspore.nn as nn
import mindspore.numpy as mnp
import mindspore.ops as ops
import numpy as np
from lstm import SentimentNet
from mindspore import Tensor, load_checkpoint, load_param_into_net
from mindspore.common.initializer import HeUniform, Uniform

mindspore.set_context(mode=mindspore.PYNATIVE_MODE)

def load_glove():
    embeddings = []
    tokens = []
    with open(r"./glove.6B.100d.txt", encoding='utf-8') as gf:
        for glove in gf:
            word, embedding = glove.split(maxsplit=1)
            tokens.append(word)
            embeddings.append(np.fromstring(embedding, dtype=np.float32, sep=' '))
    # 添加 <unk>, <pad> 两个特殊占位符对应的embedding
    embeddings.append(np.random.rand(100))
    embeddings.append(np.zeros((100,), np.float32))

    vocab = dataset.text.Vocab.from_list(tokens, special_tokens=["<unk>", "<pad>"], special_first=False)
    embeddings = np.array(embeddings).astype(np.float32)
    return vocab, embeddings


def predict_sentiment(model, vocab, sentence):
    model.set_train(False)
    trantab = str.maketrans(string.punctuation, ' '*len(string.punctuation))
    review = sentence.rstrip("\n\r").translate(trantab).lower()
    tokenized = re.split(' |,', review.lower())
    # tokenized = sentence.lower().split()
    while len(tokenized) < 400:
        tokenized.append('<pad>')
    indexed = vocab.tokens_to_ids(tokenized)
    tensor = mindspore.Tensor(indexed, mindspore.int32)
    tensor = tensor.expand_dims(0)
    prediction = model(tensor)
    prediction = mindspore.nn.Softmax(axis=-1)(prediction)
    return prediction.asnumpy()

# 加载网路


vocab, embeddings = load_glove()
pad_idx = vocab.tokens_to_ids('<pad>')
net1 = SentimentNet(embeddings, pad_idx)
net2 = SentimentNet(embeddings, pad_idx)

param_dict = load_checkpoint("./lstm_1.ckpt")
load_param_into_net(net1, param_dict)

param_dict = load_checkpoint("./lstm_2.ckpt")
load_param_into_net(net2, param_dict)

def prefict_emotion(sentence, direction):
    global net, vocab
    if direction == 1:
        
        pred = predict_sentiment(net1, vocab, sentence)[0]
        result = {
            "评分为 1 分, 概率为": pred[0].item(),
            "评分为 2 分, 概率为": pred[1].item(),
            "评分为 3 分, 概率为": pred[2].item(),
            "评分为 4 分, 概率为": pred[3].item(),
            "评分为 5 分, 概率为": pred[4].item(),
        }
        return result
    else:
        pred1 = predict_sentiment(net1, vocab, sentence)[0]
        pred2 = predict_sentiment(net2, vocab, sentence)[0]
        pred = pred1 + pred2
        pred = pred / pred.sum()
        result = {
            "评分为 1 分, 概率为": pred[0].item(),
            "评分为 2 分, 概率为": pred[1].item(),
            "评分为 3 分, 概率为": pred[2].item(),
            "评分为 4 分, 概率为": pred[3].item(),
            "评分为 5 分, 概率为": pred[4].item(),
        }
        return result

gr.Interface(
    fn=prefict_emotion,
    inputs=[gr.inputs.Textbox(
        lines=3,
        placeholder="Type a phrase that has some emotion",
        label="Input Text",
    ),gr.inputs.Radio(choices=['多模型投票', '单模型投票'],
                                            type='index')],
    outputs="label",
    title="昇思 AI 挑战赛-文本分类",
    examples=[[
        "This film is fucking terrible",
        '多模型投票',],["This film is fucking great", '单模型投票']]
).launch(share=True)

